#!/usr/bin/env python3
import argparse
import concurrent.futures
import itertools
import os
import re
import subprocess
from typing import Dict, List, Tuple


def run_tmalign(tmalign_exe: str, pdb_a: str, pdb_b: str) -> Tuple[float, float, float, str]:
    """
    运行 TM-align 对 pdb_a 与 pdb_b 进行比对。

    返回:
        tm_ab: 以 A 为参考行里出现的 TM-score（输出第一行）
        tm_ba: 以 B 为参考行里出现的 TM-score（输出第二行）
        rmsd:  RMSD (Å)
        raw:   原始输出
    说明:
        TM-align 输出中通常包含两行 TM-score：
        - 第一行: "TM-score=... (if normalized by length of Chain_1, ...)"
        - 第二行: "TM-score=... (if normalized by length of Chain_2, ...)"
        其中第二行可视为以第二个结构为参考的 TM-score。
    """
    out = subprocess.check_output([tmalign_exe, pdb_a, pdb_b], text=True)
    tms = re.findall(r"TM-score=\s*([0-9.]+)", out)
    if not tms:
        raise RuntimeError("未在 TM-align 输出中解析到 TM-score")
    if len(tms) == 1:
        tm_ab = float(tms[0])
        tm_ba = float(tms[0])
    else:
        tm_ab = float(tms[0])
        tm_ba = float(tms[1])
    rms_match = re.search(r"RMSD=\s*([0-9.]+)", out)
    rmsd = float(rms_match.group(1)) if rms_match else float("nan")
    return tm_ab, tm_ba, rmsd, out


def select_tm(tm_ab: float, tm_ba: float, policy: str) -> float:
    if policy == "second":
        return tm_ba
    if policy == "avg":
        return (tm_ab + tm_ba) / 2.0
    # 默认: 取较大值，偏保守评估相似度
    return max(tm_ab, tm_ba)


def main() -> int:
    parser = argparse.ArgumentParser(description="对目录中所有 PDB 做成对 TM-align 比较，输出矩阵与明细")
    parser.add_argument("--dir", default="/data1/yzq/good_back_RFDiffusion/test_pyg_outputs_1", help="包含 PDB 的目录")
    parser.add_argument("--suffix", default=".pdb", help="PDB 文件后缀过滤，例如 .pdb")
    parser.add_argument("--tmalign", default="/mnt/yzq_4090/rf/.local/bin/TMalign", help="TMalign 可执行文件路径")
    parser.add_argument("--tm_policy", choices=["max", "avg", "second"], default="max", help="TM-score 选取策略")
    parser.add_argument("--jobs", type=int, default=os.cpu_count() or 8, help="并行任务数")
    parser.add_argument("--save_aln", action="store_true", help="保存每个比对的原始输出到 .tmalign.log")
    parser.add_argument("--out_prefix", default="pairwise_tmalign", help="输出文件前缀")
    args = parser.parse_args()

    pdb_dir = os.path.abspath(args.dir)
    if not os.path.isdir(pdb_dir):
        print(f"[ERROR] 目录不存在: {pdb_dir}")
        return 1

    pdb_files: List[str] = [os.path.join(pdb_dir, f) for f in os.listdir(pdb_dir) if f.endswith(args.suffix)]
    pdb_files = sorted(pdb_files)
    if len(pdb_files) < 2:
        print(f"[WARN] 找到的 PDB 数量 < 2: {len(pdb_files)}")
        return 0

    names: List[str] = [os.path.basename(p) for p in pdb_files]
    index_of: Dict[str, int] = {name: i for i, name in enumerate(names)}
    n = len(names)

    tm_mat = [[float("nan") for _ in range(n)] for _ in range(n)]
    rms_mat = [[float("nan") for _ in range(n)] for _ in range(n)]
    for i in range(n):
        tm_mat[i][i] = 1.0
        rms_mat[i][i] = 0.0

    pairs = [(i, j) for i, j in itertools.combinations(range(n), 2)]

    def worker(pair: Tuple[int, int]):
        i, j = pair
        a = pdb_files[i]
        b = pdb_files[j]
        tm_ab, tm_ba, rmsd, raw = run_tmalign(args.tmalign, a, b)
        if args.save_aln:
            with open(os.path.join(pdb_dir, f"{names[i]}__VS__{names[j]}.tmalign.log"), "w") as f:
                f.write(raw)
        tm_sel = select_tm(tm_ab, tm_ba, args.tm_policy)
        return i, j, tm_ab, tm_ba, tm_sel, rmsd

    results: List[Tuple[int, int, float, float, float, float]] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.jobs) as ex:
        for res in ex.map(worker, pairs):
            results.append(res)

    # 写明细表
    tsv_path = os.path.abspath(f"{args.out_prefix}_pairs.tsv")
    with open(tsv_path, "w") as f:
        f.write("i\tj\tname_i\tname_j\tTM(i->j)\tTM(j->i)\tTM(sel)\tRMSD(A)\n")
        for i, j, tm_ab, tm_ba, tm_sel, rmsd in results:
            f.write(f"{i}\t{j}\t{names[i]}\t{names[j]}\t{tm_ab:.6f}\t{tm_ba:.6f}\t{tm_sel:.6f}\t{rmsd:.6f}\n")

    # 填矩阵（对称）
    for i, j, tm_ab, tm_ba, tm_sel, rmsd in results:
        tm_mat[i][j] = tm_sel
        tm_mat[j][i] = tm_sel
        rms_mat[i][j] = rmsd
        rms_mat[j][i] = rmsd

    # 写矩阵（带表头）
    tm_mat_path = os.path.abspath(f"{args.out_prefix}_tm_matrix.tsv")
    with open(tm_mat_path, "w") as f:
        f.write("name\t" + "\t".join(names) + "\n")
        for i in range(n):
            row = "\t".join(f"{tm_mat[i][j]:.6f}" for j in range(n))
            f.write(f"{names[i]}\t{row}\n")

    rms_mat_path = os.path.abspath(f"{args.out_prefix}_rmsd_matrix.tsv")
    with open(rms_mat_path, "w") as f:
        f.write("name\t" + "\t".join(names) + "\n")
        for i in range(n):
            row = "\t".join(f"{rms_mat[i][j]:.6f}" for j in range(n))
            f.write(f"{names[i]}\t{row}\n")

    print("完成: ")
    print(f"- 明细: {tsv_path}")
    print(f"- TM 矩阵: {tm_mat_path}")
    print(f"- RMSD 矩阵: {rms_mat_path}")
    print(f"PDB 数量: {n}，对数: {len(pairs)}，TM 选取策略: {args.tm_policy}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())


